// SPDX-FileCopyrightText: 2023-2025 Sangfor Technologies Inc.
// SPDX-License-Identifier: Mulan PSL v2
package com.geniusai.aip.auth.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.geniusai.aip.auth.config.ExternalConfig;
import com.geniusai.aip.auth.constants.LoadConstants;
import com.geniusai.aip.auth.domain.LoginEntity;
import com.geniusai.aip.auth.domain.ResponseTokenDTO;
import com.geniusai.aip.auth.domain.UserDO;
import com.geniusai.aip.auth.exception.AuthLoginException;
import com.geniusai.aip.auth.exception.AuthLoginExceptionEnum;
import com.geniusai.aip.auth.mapper.UserMapper;
import com.geniusai.aip.auth.service.LoginService;
import com.geniusai.aip.auth.service.UserService;
import com.geniusai.aip.feature.common.utils.DateUtil;
import com.geniusai.aip.feature.common.utils.json.JacksonUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.jwt.Jwt;
import org.springframework.security.jwt.crypto.sign.RsaSigner;
import org.springframework.security.jwt.crypto.sign.RsaVerifier;
import org.springframework.security.jwt.crypto.sign.SignatureVerifier;
import org.springframework.security.jwt.crypto.sign.Signer;
import org.springframework.stereotype.Service;
import org.springframework.security.jwt.JwtHelper;
import org.springframework.util.StringUtils;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureAlgorithm;

import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;

import java.security.KeyPair;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.sql.Timestamp;
import java.util.Base64;
import java.util.List;
import java.util.UUID;

import java.util.HashMap;
import java.util.Map;

/**
 * @author hml14563
 * @date 2021/3/20
 */
@Slf4j
@Service
public class LoginServiceImpl extends ServiceImpl<UserMapper, UserDO> implements LoginService {

    @Autowired
    private ExternalConfig externalConfig;

    @Autowired
    private UserService userService;

    @Autowired
    private UserMapper userMapper;

    @Autowired
    private Map<Integer, Map<String, List<String>>> rolePermission;

    @Autowired
    private KeyPair keyPair;

    private static final String LOGIN = "login";
    private static final Integer INTERVAL = 300000;
    private static final Integer COOKIE_EXPIRATION = 15 * 24 * 3600;
    private static final Long SECONDS = 1000L;
    private static final Long S3_TOKEN_TIMEOUT = 864000L;

    @Override
    public ResponseTokenDTO login(HttpServletRequest request, LoginEntity loginEntity, HttpServletResponse response) {
        UserDO userDO = userMapper.selectOne(new QueryWrapper<UserDO>().eq("username", loginEntity.getUsername()));
        if (userDO == null) {
            throw new AuthLoginException(AuthLoginExceptionEnum.USER_NOT_FOUND);
        }
        userDO.setUpdateTime(DateUtil.getDateTime());

        HttpSession session = request.getSession();

        if (!userService.queryUserByUsernameAndPassword(loginEntity.getUsername(), loginEntity.getPassword())) {
            throw new AuthLoginException(AuthLoginExceptionEnum.USER_NOT_FOUND);
        }

        ResponseTokenDTO responseTokenDTO = new ResponseTokenDTO();

        // 是ngsoc账号,直接返回token

        String jti = new String(Base64.getEncoder().encode(UUID.randomUUID().toString().getBytes()));

        String token = generateToken(externalConfig.getAccesstokenSeconds(), jti, loginEntity.getUsername());
        responseTokenDTO.setIsFirstLogin(userDO.getIsFirstLogin() == 1 ? Boolean.TRUE : Boolean.FALSE);
        responseTokenDTO.setUserId(loginEntity.getUsername());

        // 如果不是需要校验账号密码
        if (!loginEntity.getUsername().equals(externalConfig.getNgsocName()) || !loginEntity.getCode().equals(externalConfig.getNgsocCode())) {
            checkValidationCode(request, loginEntity);
            responseTokenDTO.setIsFirstLogin(userDO.getIsFirstLogin() == 1 ? Boolean.TRUE : Boolean.FALSE);
        }
        responseTokenDTO.setRole(userService.getRoleByUsername(loginEntity.getUsername()));

        String s3Access = userService.getS3AccessKeyByUsername(loginEntity.getUsername());
        String s3Secret = userService.getS3SecretKeyByUsername(loginEntity.getUsername());
        responseTokenDTO.setS3Token(getS3Token(s3Access, s3Secret));

        userMapper.updateById(userDO);
        session.removeAttribute("code");
        session.setAttribute(LOGIN, true);
        session.setMaxInactiveInterval(externalConfig.getAccesstokenSeconds());

        Cookie cookie = new Cookie("access_token", token);
        cookie.setHttpOnly(true);
        cookie.setSecure(true);
        cookie.setPath("/");
        cookie.setMaxAge(COOKIE_EXPIRATION);
        response.addCookie(cookie);

        // 将状态位设置为0，表示不是第一次登陆
        userDO.setIsFirstLogin(0);
        userMapper.updateById(userDO);
        return responseTokenDTO;
    }

    @Override
    public void logout(HttpServletRequest request, HttpSession session) {
        session.invalidate();
        log.info("the account logout");
    }

    @Override
    public String getToken(HttpServletRequest request, String username) {
        HttpSession session = request.getSession();
        if (session.getAttribute(LOGIN) == null || request.getSession(false) == null) {
            throw new AuthLoginException(AuthLoginExceptionEnum.INVALID_SESSION);
        }

        String jti = new String(Base64.getEncoder().encode(UUID.randomUUID().toString().getBytes()));

        return generateToken(externalConfig.getAccesstokenSeconds(), jti, username);
    }

    @Override
    public void getPermission(HttpServletRequest request, HttpServletResponse response) {
        String token = request.getHeader("req_access_token");
        String uri = request.getHeader("req_path");
        String method = request.getHeader("req_method");

        // 解析jwt
        RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
        SignatureVerifier verifier = new RsaVerifier(publicKey);

        Jwt jwt = JwtHelper.decodeAndVerify(token, verifier);
        String claimsStr = jwt.getClaims();
        Map claims = JacksonUtils.toObj(claimsStr, Map.class);
        String username = (String) claims.get("user_name");
        int role = userService.getRoleByUsername(username);
        boolean permit = checkUriPermission(uri, method, role, username, claims);

        String jti = new String(Base64.getEncoder().encode(UUID.randomUUID().toString().getBytes()));
        String refreshToken = generateToken(externalConfig.getAccesstokenSeconds(), jti, username);

        response.addHeader("access_token", permit ? refreshToken : "");
    }

    private boolean checkUriPermission(String uri, String method, Integer role, String username, Map claim) {
        long exp = (Integer) claim.get("exp");
        if (exp < new Timestamp(System.currentTimeMillis()).getTime() / SECONDS) {
            return false;
        }
        List<String> rules = rolePermission.get(role).get(method.toLowerCase());
        if (role.equals(LoadConstants.ADMIN)) {
            for (String rule : rules) {
                if (uri.startsWith(rule)) {
                    return true;
                }
            }
            return false;
        }
        // if not admin user
        for (String rule : rules) {
            String newRule = "";
            if (rule.contains(LoadConstants.USERINFO)) {
                newRule = rule.replace(LoadConstants.USERINFO, username);
            } else if (rule.contains(LoadConstants.RESOURCEINFO)) {
                newRule = rule.replace(LoadConstants.RESOURCEINFO, username + "_");
            }
            // prefix check
            if (uri.startsWith(newRule)) {
                return true;
            }
        }
        return false;
    }

    private String generateToken(int expireSecs, String jti, String username) {

        Map<String, Object> tokenMap = new HashMap<>(10);
        tokenMap.put("sub", "test@aip.geniusai.com");
        tokenMap.put("user_name", username);
        tokenMap.put("scope", "all");
        tokenMap.put("iss", "test@aip.geniusai.com");
        tokenMap.put("exp", new Timestamp(System.currentTimeMillis()).getTime() / SECONDS + expireSecs);
        tokenMap.put("authorities", "admin");
        tokenMap.put("jti", jti);
        tokenMap.put("client_id", "test");

        Signer signer = new RsaSigner((RSAPrivateKey) keyPair.getPrivate());

        return JwtHelper.encode(JacksonUtils.formatMap(tokenMap), signer).getEncoded();
    }

    private void checkValidationCode(HttpServletRequest request, LoginEntity loginEntity) {
        // 从session中获取正确的验证码
        HttpSession session = request.getSession();
        String code = (String) session.getAttribute("code");
        Long time = (Long) session.getAttribute("date");
        if (System.currentTimeMillis() - time > INTERVAL) {
            throw new AuthLoginException(AuthLoginExceptionEnum.VALIDATION_CODE_EXPIRED);
        }
        // 从参数中获取用户输入的验证码

        String verifyCode = loginEntity.getCode();
        if (StringUtils.isEmpty(verifyCode)) {
            // 清空session中的验证码，让用户重新获取
            throw new AuthLoginException(AuthLoginExceptionEnum.EMPTY_VALIDATION_CODE);
        }
        // 验证码校验
        if (!verifyCode.equalsIgnoreCase(code)) {
            throw new AuthLoginException(AuthLoginExceptionEnum.VALIDATION_CODE_ERROR);
        }
    }

    private String getS3Token(String accessKey, String secretKey) {
        Map<String, Object> stringObjectMap = new HashMap<>(2);
        stringObjectMap.put("alg", "HS512");
        stringObjectMap.put("typ", "JWT");
        String payload = String.format(
                "{\"accessKey\":\"%s\", \"sub\": \"minio\", \"exp\":%d}",
                accessKey, System.currentTimeMillis() / SECONDS + S3_TOKEN_TIMEOUT
        );
        String compactJws = Jwts.builder()
                .setHeader(stringObjectMap)
                .setPayload(payload)
                .signWith(SignatureAlgorithm.HS512, secretKey.getBytes())
                .compact();
        return compactJws;
    }

}
